package cn.foolishbird.crow.step;

import org.apache.commons.collections4.CollectionUtils;

import java.util.Collection;
import java.util.Objects;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;
import java.util.function.Consumer;

/**
 * @author foolish bird
 * @email luohongtu@foolishbird.cn
 */
public class IteratorExecutorDispose<E> {

    /**
     * 线程数量
     */
    private final int nThreads;

    /**
     * 线程池
     */
    private final ExecutorService executorService;

    /**
     * 支持异步读取
     */
    private final PageExecutorReadOnlyIterator<E> readOnlyIterator;

    public IteratorExecutorDispose(PageExecutorReadOnlyIterator<E> readOnlyIterator) {
        this.nThreads = 4;
        this.executorService = new ThreadPoolExecutor(nThreads, nThreads,
                0L, TimeUnit.MILLISECONDS,
                new LinkedBlockingQueue<Runnable>());
        this.readOnlyIterator = readOnlyIterator;
    }

    public IteratorExecutorDispose(int nThreads, PageExecutorReadOnlyIterator<E> readOnlyIterator) {
        this.nThreads = nThreads;
        this.executorService = new ThreadPoolExecutor(nThreads, nThreads,
                0L, TimeUnit.MILLISECONDS,
                new LinkedBlockingQueue<Runnable>());
        this.readOnlyIterator = readOnlyIterator;
    }

    public IteratorExecutorDispose(int nThreads, ExecutorService executorService, PageExecutorReadOnlyIterator<E> readOnlyIterator) {
        this.nThreads = nThreads;
        this.executorService = executorService;
        this.readOnlyIterator = readOnlyIterator;
    }

    /**
     * 分步处理
     *
     * @param consumer 执行器
     */
    public void stepDispose(Consumer<Collection<E>> consumer) {
        if (Objects.isNull(readOnlyIterator)) {
            throw new RuntimeException("步进读取器不能为空");
        }

        for (int i = 0; i < nThreads; i++) {
            int initPage = i + 1;
            executorService.execute(() -> {
                readOnlyIterator.setPageNum(initPage);

                Collection<E> collection = readOnlyIterator.read();
                while (CollectionUtils.isNotEmpty(collection)) {
                    consumer.accept(collection);
                    readOnlyIterator.forward();
                    collection = readOnlyIterator.read();
                }
            });
        }

    }
}
